from deepbots.supervisor.wrappers.tensorboard_wrapper import TensorboardLogger
import numpy as np

class TensorboardLoggerMA(TensorboardLogger):
    def __init__(self,
                 controller,
                 log_dir="logs/results",
                 v_action=0,
                 v_observation=0,
                 v_reward=0,
                 windows=[10, 100, 200]):
        super(TensorboardLoggerMA,self).__init__(controller,log_dir,v_action,v_observation,v_reward,windows)
    
    def step(self, action):
        observation, reward, isDone, info = self.controller.step(action)

        if (self.v_action > 1):
            self.file_writer.add_histogram(
                "Actions/Per Global Step",
                action,
                global_step=self.step_global)

        if (self.v_observation > 1):
            self.file_writer.add_histogram(
                "Observations/Per Global Step",
                observation,
                global_step=self.step_global)

        if (self.v_reward > 1):
            self.file_writer.add_scalar("Rewards/Per Global Step", reward,
                                        self.step_global)

        if np.all(isDone):
            self.file_writer.add_scalar(
                "Is Done/Per Reset step",
                self.step_cntr,
                global_step=self.step_reset)

        self.file_writer.flush()

        self.score += np.sum(reward)

        self.step_cntr += 1
        self.step_global += 1

        return observation, reward, isDone, info